package com.xiashitech.agent.instrumentation.db.updateset.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.constants.Instrumentation;
import com.xiashitech.agent.instrumentation.db.common.XSUpdateSizeContext;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSUpdateStatementCallDepth;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSUpdateStatementUpdateCountCallDepth;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.sql.PreparedStatement;
import java.sql.Statement;
import java.util.Arrays;

import static com.xiashitech.agent.instrumentation.db.updateset.XSUpdateSetSingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSUpdateStatementInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("java.sql.Statement");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return implementsInterface(named("java.sql.Statement"));
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                named("executeUpdate")
                        .and(takesArgument(0,String.class))
                        .and(returns(int.class))
                        .and(isPublic())
                        .and(not(isFinal())),
                XSUpdateStatementInstrumentation.class.getName() + "$XSExecuteUpdateAdvice");
        transformer.applyAdviceToMethod(
                named("executeUpdate")
                        .and(takesArgument(0,String.class))
                        .and(returns(int.class))
                        .and(isPublic())
                        .and(isFinal()),
                XSUpdateStatementInstrumentation.class.getName() + "$XSExecuteUpdateAdvice");
        transformer.applyAdviceToMethod(
                named("execute")
                        .and(takesArgument(0, String.class))
                        .and(returns(boolean.class))
                        .and(isPublic())
                        .and(not(isFinal())),
                XSUpdateStatementInstrumentation.class.getName() + "$XSExecuteAdvice");
        transformer.applyAdviceToMethod(
                named("execute")
                        .and(takesArgument(0, String.class))
                        .and(returns(boolean.class))
                        .and(isPublic())
                        .and(isFinal()),
                XSUpdateStatementInstrumentation.class.getName() + "$XSExecuteAdvice");
        transformer.applyAdviceToMethod(
                named("getUpdateCount")
                        .and(takesArguments(0))
                        .and(returns(int.class))
                        .and(isPublic())
                        .and(not(isFinal())),
                XSUpdateStatementInstrumentation.class.getName() + "$XSUpdateCountAdvice");
        transformer.applyAdviceToMethod(
                named("getUpdateCount")
                        .and(takesArguments(0))
                        .and(returns(int.class))
                        .and(isPublic())
                        .and(isFinal()),
                XSUpdateStatementInstrumentation.class.getName() + "$XSUpdateCountAdvice");
        transformer.applyAdviceToMethod(
                named("addBatch")
                        .and(takesArgument(0, String.class))
                        .and(returns(void.class))
                        .and(isPublic())
                        .and(not(isFinal())),
                XSUpdateStatementInstrumentation.class.getName() + "$XSAddBatchAdvice");
        transformer.applyAdviceToMethod(
                named("addBatch")
                        .and(takesArgument(0, String.class))
                        .and(returns(void.class))
                        .and(isPublic())
                        .and(isFinal()),
                XSUpdateStatementInstrumentation.class.getName() + "$XSAddBatchAdvice");
        transformer.applyAdviceToMethod(
                named("executeBatch")
                        .and(takesArguments(0))
                        .and(returns(int[].class))
                        .and(isPublic()),
                XSUpdateStatementInstrumentation.class.getName() + "$XSExecuteBatchAdvice");
    }

    /**
     * 处理executeUpdate插装拦截
     */
    @SuppressWarnings("unused")
    public static class XSExecuteUpdateAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnterForResultSetPrepare(
                @Advice.This Statement statement,
                @Advice.Argument(0) String sql,
                @Advice.Local("otelRequest") XSUpdateSizeContext.Request request,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSUpdateStatementCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            Context parentContext = currentContext();
            request = XSUpdateSizeContext.Request.builder().sql(sql).build();
            if (request == null || !instrumenter().shouldStart(parentContext, request)) {
                return;
            }
            parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
            context = instrumenter().start(parentContext, request);
            context = BaggageUtil.syncBaggage().build().storeInContext(context);
            scope = context.makeCurrent();
            durationPrinter.printElapseDuration("XSExecuteBatchAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExitForResultSetStart(
                @Advice.This Statement statement,
                @Advice.Local("otelRequest") XSUpdateSizeContext.Request request,
                @Advice.Thrown Throwable throwable,
                @Advice.Return int updateSize,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            request.setSize(updateSize);
            if (scope != null) {
                scope.close();
                instrumenter().end(context, request, null, throwable);
            }
            durationPrinter.printElapseDuration("XSExecuteBatchAdvice e");
        }
    }

    /**
     * 处理execute插装拦截-开始
     */
    @SuppressWarnings("unused")
    public static class XSExecuteAdvice {
        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.This Statement statement,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSUpdateStatementCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            durationPrinter.printElapseDuration("XSExecuteAdvice b");
        }
        /**
         * 执行完execute方法就开始span
         * @param statement
         * @param throwable
         * @param executeType
         * @param callDepth
         */
        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.This Statement statement,
                @Advice.Argument(0) String sql,
                @Advice.Thrown Throwable throwable,
                @Advice.Return boolean executeType,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            if(throwable == null && !executeType) {
                XSUpdateSizeContext.refreshCurrentThreadLocal();
                XSUpdateSizeContext.setStatementWithContext(statement, XSUpdateSizeContext.UpdateSizeContext.builder().request(XSUpdateSizeContext.Request.builder().sql(sql).build()).executeType(executeType).build());
            }
            durationPrinter.printElapseDuration("XSExecuteAdvice e");
        }
    }

    /**
     * 处理execute插装拦截-结束
     */
    @SuppressWarnings("unused")
    public static class XSUpdateCountAdvice {
        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.This Statement statement,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSUpdateStatementUpdateCountCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            //由于PreparedStatement继承Statement的方法，所以防止生成两边该span
            XSUpdateSizeContext.UpdateSizeContext preparedUpdateSizeContext = XSUpdateSizeContext.getUpdateSizeContext(statement);
            if(preparedUpdateSizeContext != null && statement instanceof PreparedStatement) {
                return;
            }
            durationPrinter.printElapseDuration("XSUpdateCountAdvice b");
        }
        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.This Statement statement,
                @Advice.Thrown Throwable throwable,
                @Advice.Return int updateSize,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            XSUpdateSizeContext.UpdateSizeContext updateSizeContext = XSUpdateSizeContext.getUpdateSizeContext(statement);
            if(updateSizeContext == null || updateSizeContext.isExecuteType() || (updateSizeContext != null && statement instanceof PreparedStatement)) {//由于PreparedStatement继承Statement的方法，所以最后括号里的条件防止生成两遍该span
                return;
            }
            Context parentContext = currentContext();
            XSUpdateSizeContext.Request request = updateSizeContext.getRequest();
            if (request == null || !instrumenter().shouldStart(parentContext, request)) {
                return;
            }
            //生成span
            if(updateSize != -1) {
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                Context context = instrumenter().start(parentContext, request);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                request.setSize(updateSize);
                instrumenter().end(context, request, null, throwable);
            }
            durationPrinter.printElapseDuration("XSUpdateCountAdvice e");
        }
    }

    /**
     * 处理addBatch插装拦截
     */
    @SuppressWarnings("unused")
    public static class XSAddBatchAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnterForResultSetPrepare(
                @Advice.This Statement statement,
                @Advice.Argument(0) String sql,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSUpdateStatementCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            XSUpdateSizeContext.UpdateSizeContext updateSizeContext = XSUpdateSizeContext.getUpdateSizeContext(statement);
            if(updateSizeContext == null) {
                XSUpdateSizeContext.setStatementWithContext(statement, XSUpdateSizeContext.UpdateSizeContext.builder().request(XSUpdateSizeContext.Request.builder().sql(sql).build()).build());
            } else {
                XSUpdateSizeContext.Request request = updateSizeContext.getRequest();
                if(request == null || request.getSql() == null || "".equals(request.getSql())) {
                    XSUpdateSizeContext.setStatementWithContext(statement, XSUpdateSizeContext.UpdateSizeContext.builder().request(XSUpdateSizeContext.Request.builder().sql(sql).build()).build());
                } else {
                    request.setSql(request.getSql() + Instrumentation.SEMICOLON + sql);
                }
            }
            updateSizeContext.setFromAddBatch(true);
            durationPrinter.printElapseDuration("XSAddBatchAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExitForResultSetStart(
                @Advice.This Statement statement,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            durationPrinter.printElapseDuration("XSAddBatchAdvice e");
        }
    }

    /**
     * 处理executeBatch插装拦截
     */
    @SuppressWarnings("unused")
    public static class XSExecuteBatchAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnterForResultSetPrepare(
                @Advice.This Statement statement,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSUpdateStatementCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            XSUpdateSizeContext.UpdateSizeContext updateSizeContext = XSUpdateSizeContext.getUpdateSizeContext(statement);
            if(updateSizeContext == null || !updateSizeContext.isFromAddBatch()) {
                return;
            }
            durationPrinter.printElapseDuration("XSExecuteBatchAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExitForResultSetStart(
                @Advice.This Statement statement,
                @Advice.Thrown Throwable throwable,
                @Advice.Return int[] updateSize,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSUpdateStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            XSUpdateSizeContext.UpdateSizeContext updateSizeContext = XSUpdateSizeContext.getUpdateSizeContext(statement);
            if(updateSizeContext == null || !updateSizeContext.isFromAddBatch()) {
                return;
            }

            Context parentContext = currentContext();
            XSUpdateSizeContext.Request request = updateSizeContext.getRequest();
            if (request == null || !instrumenter().shouldStart(parentContext, request)) {
                return;
            }
            parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
            //生成span
            Context context = instrumenter().start(parentContext, request);
            context = BaggageUtil.syncBaggage().build().storeInContext(context);
            request.setSize(Arrays.stream(updateSize).sum());
            instrumenter().end(context, request, null, throwable);
            XSUpdateSizeContext.refreshCurrentThreadLocal();
            durationPrinter.printElapseDuration("XSExecuteBatchAdvice e");
        }
    }
}
